#include "analyzer.h"
#include "debug_utils.h"

mypointsto::disjoint_set mypointsto::edge_record::s;

mypointsto::edge_record::edge_record(int s, int d){
	v_src = s;
	v_dst = d;
}

bool mypointsto::edge_record::operator < (const edge_record &rhs) const{
	return (v_src < rhs.v_src || (v_src == rhs.v_src && v_dst < rhs.v_dst));
}

mypointsto::strong_component::strong_component(){
	no_deref_node = true;
	non_deref_v = -1;
}

mypointsto::strong_component::strong_component(bool b, int v){
	no_deref_node = b;
	non_deref_v = v;
}

mypointsto::analyzer::analyzer(){
	this->cons_builder = NULL;
}

mypointsto::analyzer::analyzer(mypointsto::constraint_builder *cb){
	this->cons_builder = cb;
}

int mypointsto::analyzer::initialize(){
	init_graph();
	scc.resize((int)(cons_builder->labels.size()));//must make sure resize method is only invoked here!
	return 0;
}

int mypointsto::analyzer::initialize_hybrid(){
	init_hybrid_graph();
	scc.resize((int)(cons_builder->labels.size()+deref_map.size()));
	return 0;
}

int mypointsto::analyzer::init_graph(){
	if(!cons_builder)
		return -1;
	int sz = (int) (cons_builder->labels.size());
	sz = mypointsto::round_up(sz, 32);
	for(int i=0; i<(int)((cons_builder->labels).size()); i++){
		points_to pt_info(boost::dynamic_bitset<unsigned int> (32));
		load_from ld_info(boost::dynamic_bitset<unsigned int>(32), pt_info);
		store_to st_info(boost::dynamic_bitset<unsigned int> (32), ld_info);
		component_index comp_idx(i, st_info);
		vertex_info v_info(i, comp_idx);
		wks.push(i);
		boost::add_vertex(v_info, this->graph);
	}
	points_to_map ptm = boost::get(points_to_t(), this->graph);
	store_to_map stm = boost::get(store_to_t(), this->graph);
	load_from_map lfm = boost::get(load_from_t(), this->graph);
	index_map indices = boost::get(boost::vertex_index_t(), this->graph);
	vertices_range vp = boost::vertices(this->graph);
	for(; vp.first!=vp.second; vp.first++){
		ptm[*(vp.first)].resize(sz);
		stm[*(vp.first)].resize(sz);
		lfm[*(vp.first)].resize(sz);
	}
	base_constraint *bc = NULL;
	simple_constraint *sc = NULL;
	load_constraint *lc = NULL;
	store_constraint *stc = NULL;
	for(std::vector<constraint*>::iterator it = (cons_builder->constraints).begin(); it != (cons_builder->constraints).end(); it++){
		switch ((*it)->type){
			case base_t:
				bc = (base_constraint*) (*it);
				process_base_constraint(bc->ptr_id, bc->val_id, true);
				break;
			case simple_t:
				sc = (simple_constraint*) (*it);
				process_simple_constraint(sc->src_id, sc->dst_id, true);
				break;
			case load_t:
				lc = (load_constraint*) (*it);
				process_load_constraint(lc->addr_id, lc->val_id, true);
				break;
			case store_t:
				stc = (store_constraint*)(*it);
				process_store_constraint(stc->val_id, stc->addr_id, true);
			default:
				break;
		}
	}
	return 0;
}

int mypointsto::analyzer::init_hybrid_graph(){
	if(!cons_builder)
		return -1;
	int sz = (int) (cons_builder->labels.size());
	sz = mypointsto::round_up(sz, 32);
	for(int i=0; i<(int)((cons_builder->labels).size()); i++){
		points_to pt_info(boost::dynamic_bitset<unsigned int> (32));
		load_from ld_info(boost::dynamic_bitset<unsigned int>(32), pt_info);
		store_to st_info(boost::dynamic_bitset<unsigned int> (32), ld_info);
		component_index comp_idx(i, st_info);
		vertex_info v_info(i, comp_idx);
		wks.push(i);
		boost::add_vertex(v_info, this->graph);
	}
	points_to_map ptm = boost::get(points_to_t(), this->graph);
	store_to_map stm = boost::get(store_to_t(), this->graph);
	load_from_map lfm = boost::get(load_from_t(), this->graph);
	index_map indices = boost::get(boost::vertex_index_t(), this->graph);
	vertices_range vp = boost::vertices(this->graph);
	for(; vp.first!=vp.second; vp.first++){
		ptm[*(vp.first)].resize(sz);
		stm[*(vp.first)].resize(sz);
		lfm[*(vp.first)].resize(sz);
	}
	base_constraint *bc = NULL;
	simple_constraint *sc = NULL;
	load_constraint *lc = NULL;
	store_constraint *stc = NULL;
	for(std::vector<constraint*>::iterator it = (cons_builder->constraints).begin(); it != (cons_builder->constraints).end(); it++){
		switch ((*it)->type){
			case base_t:
				bc = (base_constraint*) (*it);
				process_base_constraint(bc->ptr_id, bc->val_id, true);
				break;
			case simple_t:
				sc = (simple_constraint*) (*it);
				process_simple_constraint(sc->src_id, sc->dst_id, true);
				break;
			case load_t:
				lc = (load_constraint*) (*it);
				process_load_constraint(lc->addr_id, lc->val_id, true);
				add_deref_vertex(lc->addr_id, lc->val_id, load_t);
				break;
			case store_t:
				stc = (store_constraint*)(*it);
				process_store_constraint(stc->val_id, stc->addr_id, true);
				add_deref_vertex(stc->val_id, stc->addr_id, store_t);
			default:
				break;
		}
	}
	return 0;
}

void mypointsto::analyzer::analyze(bool lazy){
	index_map indices = boost::get(boost::vertex_index_t(), this->graph);
	store_to_map stm = boost::get(store_to_t(), this->graph);
	load_from_map lfm = boost::get(load_from_t(), this->graph);
	points_to_map ptm = boost::get(points_to_t(), this->graph);
	if(lazy)
		cycle_detect_edges.clear();
	while(!wks.empty()){
		int cur_id = wks.top();
		wks.pop();
		vertices_range vp1;
		if(find_vertex_by_id(this->graph, scc.find(cur_id), vp1)){
			int sz1 = (int) ptm[*(vp1.first)].size();
			for(int i=0; i<sz1; i++){
				if(!ptm[*(vp1.first)].test(i))
					continue;
				int sz2 = (int) lfm[*(vp1.first)].size();
				for(int j=0; j<sz2; j++){
					if(!lfm[*(vp1.first)].test(j))
						continue;
					vertices_range vp2, vp3;
					if(find_vertex_by_id(this->graph, scc.find(i), vp2) && find_vertex_by_id(this->graph, scc.find(j), vp3)){
						edge_search_result result = boost::edge(*(vp2.first), *(vp3.first), this->graph);
						if(!result.second && *(vp2.first)!=*(vp3.first)){
							boost::add_edge(*(vp2.first), *(vp3.first), this->graph);
							wks.push(scc.find(i));
						}
					}
				}
				int sz3 = (int) stm[*(vp1.first)].size();
				for(int j=0; j<sz3; j++){
					if(!stm[*(vp1.first)].test(j))
						continue;
					vertices_range vp2, vp3;
					if(find_vertex_by_id(this->graph, scc.find(i), vp2) && find_vertex_by_id(this->graph, scc.find(j), vp3)){
						edge_search_result result = boost::edge(*(vp3.first), *(vp2.first), this->graph);
						if(!result.second && *(vp3.first) != *(vp2.first)){
							boost::add_edge(*(vp3.first), *(vp2.first), this->graph);
							wks.push(scc.find(j));
						}
					}
				}
			}
			adj_vertices_range range = boost::adjacent_vertices(*(vp1.first), this->graph);
			for(; range.first != range.second; range.first++){
				bool flag = false;
				if(lazy){
					edge_record temp(scc.find(indices[*(vp1.first)]), scc.find(indices[*(range.first)]));
					if(cycle_detect_edges.find(temp)==cycle_detect_edges.end() && ptm[*(vp1.first)]==ptm[*(range.first)] && ptm[*(vp1.first)].count()>0){//trigger cycle detection
						dfs_path_recorder tracker;
						std::vector<Vertex> path;
						boost::depth_first_search(this->graph, boost::visitor(tracker).root_vertex(*(range.first)));//we have edge n->z, try to find out if there is a path from z to n, thus find out if there is a cycle
						flag = is_there_path(*(range.first), *(vp1.first), dfs_path_recorder::predecessor, this->graph, path);
						if(flag){
							shrink_cycle(path, this->graph, scc);
						}
					}
				}
				if(!flag){
					boost::dynamic_bitset<unsigned int> orig(ptm[*(range.first)]);
					ptm[*(range.first)] |= ptm[*(vp1.first)];
					if(orig != ptm[*(range.first)])
						wks.push(scc.find(indices[*(range.first)]));
				}
			}
		}
	}
}

void mypointsto::analyzer::analyze_hybrid(){
	//by this time, hybrid graph is built
	//compute scc, offline phase
	component_map comp_map;
	index_map indices = boost::get(boost::vertex_index_t(), this->graph);
	store_to_map stm = boost::get(store_to_t(), this->graph);
	load_from_map lfm = boost::get(load_from_t(), this->graph);
	points_to_map ptm = boost::get(points_to_t(), this->graph);
	boost::strong_components(this->graph, comp_map);
	vertices_range vp = boost::vertices(this->graph);
	std::set<int> strong_component_id;
	for(;vp.first!=vp.second; vp.first++){
		strong_component_id.insert(comp_map[*(vp.first)]);
	}
	int size = (int) strong_component_id.size();
	std::vector<strong_component> strong_comps(size);
	vp = boost::vertices(this->graph);
	for(;vp.first!=vp.second; vp.first++){
		int comp_id = comp_map[*(vp.first)];
		strong_comps[comp_id].members.insert(indices[*(vp.first)]);
		if(is_deref_vertex(indices[*(vp.first)]))
			strong_comps[comp_id].no_deref_node = false;
		if(strong_comps[comp_id].non_deref_v<0)
			strong_comps[comp_id].non_deref_v = indices[*(vp.first)];
	}
	for(int i=0; i<(int)strong_comps.size(); i++){
		if(strong_comps[i].no_deref_node && strong_comps[i].members.size()>1){
			//collapse scc
			vertices_range vp1;
			int representative = scc.find(strong_comps[i].non_deref_v);
			find_vertex_by_id(this->graph, representative, vp1);
			for(std::set<int>::iterator it=strong_comps[i].members.begin(); it!=strong_comps[i].members.end(); it++)
			{
				vertices_range vp2;
				find_vertex_by_id(this->graph, scc.find(*it), vp2);
				collapse(*(vp1.first), *(vp2.first), this->graph, scc);
			}
		}else if(!strong_comps[i].no_deref_node && strong_comps[i].members.size()>1){
			for(std::set<int>::iterator it=strong_comps[i].members.begin(); it!=strong_comps[i].members.end(); it++){
				if(*it!=strong_comps[i].non_deref_v && is_deref_vertex(*it))
					L.push_back(std::pair<int, int> (reverse_deref_map[*it], strong_comps[i].non_deref_v));
			}
		}
	}
	//online phase
	int sz = (int) cons_builder->labels.size();
	bool flag;
	/*token is used to mark if we have done anything to change the worklist during last iteration. If no and cur_id==last_id, we can just pop the worklist and restart the loop without having to perform the tedious work below. token and last_id greatly decrease running time*/
	bool token = false;
	int last_id = -1;
	while(!wks.empty()){
		token = false;
		int cur_id = wks.top();
		wks.pop();
		if(last_id == cur_id && !token)
			continue;
		last_id = cur_id;
		vertices_range vp1;
		flag = find_vertex_by_id(this->graph, scc.find(cur_id), vp1);
		for(std::vector<std::pair<int, int> >::iterator iter = L.begin(); iter != L.end(); iter++){
			if(scc.find(iter->first)==scc.find(cur_id)){
				vertices_range vp2;
				flag = find_vertex_by_id(this->graph, scc.find(iter->second), vp2);
				//int sz1 = (int) ptm[*(vp1.first)].size();
				for(int i=0; i<sz; i++){
					if(!ptm[*(vp1.first)].test(i))
						continue;
					vertices_range vp3;
					flag = find_vertex_by_id(this->graph, scc.find(i), vp3);
					if(scc.find(i)!=scc.find(indices[*(vp2.first)])){//hasn't collapsed yet, then collapse
					    collapse(*(vp2.first), *(vp3.first), this->graph, scc);
					    wks.push(scc.find(iter->second));
						token = true;
					}
				}
			}
		}

		flag = find_vertex_by_id(this->graph, scc.find(cur_id), vp1);
		//int sz1 = (int) ptm[*(vp1.first)].size();
		for(int i=0; i<sz; i++){
			if(!ptm[*(vp1.first)].test(i))
					continue;
			//int sz2 = (int) lfm[*(vp1.first)].size();
			for(int j=0; j<sz; j++){
				if(!lfm[*(vp1.first)].test(j))
					continue;
				vertices_range vp2, vp3;
				if(find_vertex_by_id(this->graph, scc.find(i), vp2) && find_vertex_by_id(this->graph, scc.find(j), vp3)){
					edge_search_result result = boost::edge(*(vp2.first), *(vp3.first), this->graph);
					if(!result.second && *(vp2.first)!=*(vp3.first)){
						boost::add_edge(*(vp2.first), *(vp3.first), this->graph);
						wks.push(scc.find(i));
						token = true;
					}
				}
			}
			//int sz3 = (int) stm[*(vp1.first)].size();
			for(int j=0; j<sz; j++){
				if(!stm[*(vp1.first)].test(j))
					continue;
				vertices_range vp2, vp3;
				if(find_vertex_by_id(this->graph, scc.find(i), vp2) && find_vertex_by_id(this->graph, scc.find(j), vp3)){
					edge_search_result result = boost::edge(*(vp3.first), *(vp2.first), this->graph);
					if(!result.second && *(vp3.first) != *(vp2.first)){
						boost::add_edge(*(vp3.first), *(vp2.first), this->graph);
						wks.push(scc.find(j));
						token = true;
					}
				}
			}
		}

		adj_vertices_range range = boost::adjacent_vertices(*(vp1.first), this->graph);
		for(; range.first != range.second; range.first++){
			boost::dynamic_bitset<unsigned int> orig(ptm[*(range.first)]);
			ptm[*(range.first)] |= ptm[*(vp1.first)];
			if(orig != ptm[*(range.first)]){
				wks.push(scc.find(indices[*(range.first)]));
				token = true;
			}
		}
	}
}

void mypointsto::analyzer::process_base_constraint(int ptr_id, int val_id, bool init) {
	vertices_range range;
	if(find_vertex_by_id(this->graph, ptr_id, range)){
		points_to_map ptm = get(points_to_t(), this->graph);
		ptm[(*range.first)].set(val_id);
	}
}

void mypointsto::analyzer::process_simple_constraint(int src, int dst, bool init){
	vertices_range vp1, vp2;
	if(!find_vertex_by_id(this->graph, src, vp1))
		return;
	if(!find_vertex_by_id(this->graph, dst, vp2))
		return;
	edge_search_result result = boost::edge(*(vp1.first), *(vp2.first), this->graph);
	if(!result.second)
		boost::add_edge(*(vp1.first), *(vp2.first), this->graph);
}

void mypointsto::analyzer::process_load_constraint(int addr_id, int val_id, bool init){
	vertices_range vp1, vp2;
	if(!find_vertex_by_id(this->graph, addr_id, vp1))
		return;
	if(!find_vertex_by_id(this->graph, val_id, vp2))
		return;
	if(init){
		index_map indices = boost::get(boost::vertex_index_t(), this->graph);
		load_from_map lfm = boost::get(load_from_t(), this->graph);
		lfm[*(vp1.first)].set(indices[*(vp2.first)]);
	}
}

void mypointsto::analyzer::process_store_constraint(int val_id, int addr_id, bool init){
	vertices_range vp1, vp2;
	if(!find_vertex_by_id(this->graph, addr_id, vp1))
		return;
	if(!find_vertex_by_id(this->graph, val_id, vp2))
		return;
	if(init){
		index_map indices = boost::get(boost::vertex_index_t(), this->graph);
		store_to_map stm = boost::get(store_to_t(), this->graph);
		stm[*(vp1.first)].set(indices[*(vp2.first)]);
	}
}

void mypointsto::analyzer::add_deref_vertex(int src, int dst, mypointsto::cons_t type){
	int size = (int) (cons_builder->labels.size()+deref_map.size());
	int sz = (int) cons_builder->labels.size();
	vertices_range vp1, vp2;
	std::map<int, int>::iterator iter;
	store_to_map stm = boost::get(store_to_t(), this->graph);
	load_from_map lfm = boost::get(load_from_t(), this->graph);
	points_to_map ptm = boost::get(points_to_t(), this->graph);
	Vertex v;
	switch (type){
		case load_t:
			if(deref_map.count(src)==0){
				deref_map.insert(std::pair<int, int> (src, size));
				reverse_deref_map.insert(std::pair<int, int> (size, src));
				points_to pt_info(boost::dynamic_bitset<unsigned int> (32));
				load_from ld_info(boost::dynamic_bitset<unsigned int>(32), pt_info);
				store_to st_info(boost::dynamic_bitset<unsigned int> (32), ld_info);
				component_index comp_idx(size, st_info);
				vertex_info v_info(size, comp_idx);
				v = boost::add_vertex(v_info, this->graph);
				lfm[v].resize(sz);
				stm[v].resize(sz);
				ptm[v].resize(sz);
			}
			iter = deref_map.find(src);
			find_vertex_by_id(this->graph, iter->second, vp1);
			find_vertex_by_id(this->graph, dst, vp2);
			boost::add_edge(*(vp1.first), *(vp2.first), this->graph);
			break;
		case store_t:
			if(deref_map.count(dst)==0){
				deref_map.insert(std::pair<int, int> (dst, size));
				reverse_deref_map.insert(std::pair<int, int> (size, dst));
				points_to pt_info(boost::dynamic_bitset<unsigned int> (32));
				load_from ld_info(boost::dynamic_bitset<unsigned int>(32), pt_info);
				store_to st_info(boost::dynamic_bitset<unsigned int> (32), ld_info);
				component_index comp_idx(size, st_info);
				vertex_info v_info(size, comp_idx);
				v = boost::add_vertex(v_info, this->graph);
				lfm[v].resize(sz);
				stm[v].resize(sz);
				ptm[v].resize(sz);
			}
			iter = deref_map.find(dst);
			find_vertex_by_id(this->graph, iter->second, vp2);
			find_vertex_by_id(this->graph, src, vp1);
			boost::add_edge(*(vp1.first), *(vp2.first), this->graph);
			break;
		default: break;
	}
}

void mypointsto::analyzer::dump_graph(std::ostream &out){
	display_graph(this->graph, out);
}

void mypointsto::analyzer::dump_analysis_result(std::ostream &out){
	int size = (int) (cons_builder->labels.size());
	points_to_map ptm = boost::get(points_to_t(), this->graph);
	for(int i=0; i<size; i++){
		int vertex_id = scc.find(i);
		vertices_range vp;
		if(find_vertex_by_id(this->graph, vertex_id, vp)){
			if(ptm[*(vp.first)].count()<=0)
				continue;
			out<<"v"<<i<<": ";
			for(int j = 0; j<(int)ptm[*(vp.first)].size(); j++){
				if(ptm[*(vp.first)].test(j))
					out<<"v"<<j<<" ";
			}
			out<<std::endl;
		}
	}
}

bool mypointsto::analyzer::is_deref_vertex(int id){
	return reverse_deref_map.find(id)!=reverse_deref_map.end();
}

bool mypointsto::find_vertex_by_id(Graph &g, int id, vertices_range &range){
	index_map indices = get(boost::vertex_index_t(), g);
	for(range=boost::vertices(g); range.first!=range.second; range.first++){
		if(indices[*(range.first)]==id){
			return true;
		}
	}
	return false;
}

void mypointsto::display_graph(Graph &g, std::ostream &out){
	vertices_range vp = boost::vertices(g);
	index_map indices = get(boost::vertex_index_t(), g);
	points_to_map ptm = get(points_to_t(), g);
	for(;vp.first!=vp.second; vp.first++){
		out<<indices[*(vp.first)]<<": ";
		int sz = (int) ptm[*(vp.first)].size();
		for(int i=0; i<sz; i++){
			if(ptm[*(vp.first)].test(i)){
				out<<i<<" ";
			}
		}
		out<<std::endl;
	}
	edges_range ep = boost::edges(g);
	for(; ep.first!=ep.second; ep.first++){
		out<<indices[boost::source(*(ep.first), g)]<<"---->"<<indices[boost::target(*(ep.first), g)]<<std::endl;
	}
}

void mypointsto::shrink_cycle(std::vector<Vertex> &path, Graph &g, mypointsto::disjoint_set &s){
	index_map indices = boost::get(boost::vertex_index_t(), g);
	points_to_map ptm = boost::get(points_to_t(), g);
	load_from_map lfm = boost::get(load_from_t(), g);
	store_to_map stm = boost::get(store_to_t(), g);

	std::vector<Edge> new_edges;
	for(int i=1; i!=(int)path.size(); i++){
		collapse(path[0], path[i], g, s);
	}
}

bool mypointsto::is_there_path(Vertex v1, Vertex v2, std::map<Vertex, Vertex> &predecessor, Graph &g, std::vector<Vertex> &path){
	if(predecessor.find(v2)==predecessor.end())
		return false;
	Vertex orig;
	orig = v2;
	while(v2!=v1){
		path.push_back(v2);
		std::map<Vertex, Vertex>::iterator it = predecessor.find(v2);
		if(it==predecessor.end())
			return false;
		v2 = it->second;
		if(orig==v2){
			return false;
		}
	}
	path.push_back(v1);
	return true;
}

void mypointsto::print_path(std::vector<Vertex> &path, Graph &g, std::ostream &out){
	index_map indices = boost::get(boost::vertex_index_t(), g);
	std::vector<Vertex>::reverse_iterator rit = path.rbegin();
	out<<"path: "<<indices[*rit];
	rit++;
	for(; rit != path.rend(); rit ++)
		out<<"->"<<indices[*rit];
	out<<std::endl;
}

int mypointsto::round_up(int n, int base){
	return n+base-n%base;
}

void mypointsto::collapse(Vertex v1, Vertex v2, Graph &g, disjoint_set &s){
	if(v1==v2)//there is no need to collapse a single vertex
		return;
	index_map indices = boost::get(boost::vertex_index_t(), g);
	points_to_map ptm = boost::get(points_to_t(), g);
	store_to_map stm = boost::get(store_to_t(), g);
	load_from_map lfm = boost::get(load_from_t(), g);
	//union bitset structures
	ptm[v1] |= ptm[v2];
	lfm[v1] |= lfm[v2];
	stm[v1] |= stm[v2];
	s.merge(indices[v1], indices[v2]);
	out_edge_range ep1 = boost::out_edges(v2, g);
	in_edge_range ep2 = boost::in_edges(v2, g);
	for(; ep1.first!=ep1.second; ep1.first++){
		Vertex t = boost::target(*(ep1.first), g);
		if(t!=v1){
			add_edge(v1, t, g);
		}
	}
	for(; ep2.first!=ep2.second; ep2.first++){
		Vertex s = boost::source(*(ep2.first), g);
		if(s!=v1){
			add_edge(s, v1, g);
		}
	}
	boost::clear_vertex(v2, g);
	boost::remove_vertex(v2, g);
}
